"""堆"""
from .exceptions import IndexErrorOfStructure, DataStructureIsEmpty


class MaxHeap:
    """
    @add   O(log(n))
    @extract_max   O(log(n))
    @replace   O(log(n))
    @heapify   O(n)
    """
    def __init__(self):
        self._data = []

    def __len__(self):
        return len(self._data)

    def is_empty(self):
        return len(self._data) == 0

    @staticmethod
    def _parent(index):
        """返回一个索引位置的父节点的索引位置"""
        if index == 0:
            raise IndexErrorOfStructure("Index-0 has no parent.")
        return (index - 1) // 2

    @staticmethod
    def _left(index):
        """返回一个索引的左孩子的索引"""
        return index * 2 + 1

    @staticmethod
    def _right(index):
        """返回一个索引的右孩子的索引"""
        return index * 2 + 2

    def add(self, e):
        """向最大堆中添加元素e"""
        self._data.append(e)
        self.shift_up(len(self) - 1)

    def shift_up(self, index):
        """最大堆元素上浮，保证每个节点的值大于其孩子节点的值"""
        while index > 0 and self._data[self._parent(index)] < self._data[index]:
            self._data[self._parent(index)], self._data[index] = self._data[index], self._data[self._parent(index)]
            index = self._parent(index)

    def find_max(self):
        """查看当前堆最大元素"""
        if self.is_empty():
            raise DataStructureIsEmpty("MaxHeap is Empty.")
        return self._data[0]

    def extract_max(self):
        """返回并删除堆中最大元素"""
        result = self.find_max()
        self._data[0], self._data[-1] = self._data[-1], self._data[0]
        self._data.pop()
        self.shift_down(0)
        return result

    def shift_down(self, index):
        """元素下沉，将较小元素放置到底层"""
        while self._left(index) < len(self):
            j = self._left(index)
            if j + 1 < len(self) and self._data[j + 1] > self._data[j]:
                j = self._right(index)
                # j为左右孩子中元素较大的节点的索引
            if self._data[index] >= self._data[j]:
                # 比左右孩子节点都大，不用下沉
                break
            self._data[index], self._data[j] = self._data[j], self._data[index]
            index = j

    def replace(self, e):
        """取出最大元素，并替换成新元素e"""
        result = self.find_max()
        self._data[0] = e
        self.shift_down(0)
        return result

    @classmethod
    def heapify(cls, arr):
        """将任意数组整理成最大堆，并返回最大堆"""
        array = list(arr)
        heap = MaxHeap()
        heap._data = array
        i = heap._parent(len(heap))
        while i >= 0:
            heap.shift_down(i)
            i -= 1
        return heap


class MinHeap:
    """
    @add   O(log(n))
    @extract_max   O(log(n))
    @replace   O(log(n))
    @heapify   o(n)
    """
    def __init__(self):
        self._data = []

    def size(self):
        return len(self._data)

    def __len__(self):
        return len(self._data)

    def is_empty(self):
        return len(self._data) == 0

    @staticmethod
    def _parent(index):
        """返回一个索引位置的父节点的索引位置"""
        if index == 0:
            raise IndexErrorOfStructure("Index-0 has no parent.")
        return (index - 1) // 2

    @staticmethod
    def _left(index):
        """返回一个索引的左孩子的索引"""
        return index * 2 + 1

    @staticmethod
    def _right(index):
        """返回一个索引的右孩子的索引"""
        return index * 2 + 2

    def add(self, e):
        """向最大堆中添加元素e"""
        self._data.append(e)
        self.shift_up(self.size() - 1)

    def shift_up(self, index):
        """最小堆元素上浮，保证每个节点的值小于其孩子节点的值"""
        while index > 0 and self._data[self._parent(index)] > self._data[index]:
            self._data[self._parent(index)], self._data[index] = self._data[index], self._data[self._parent(index)]
            index = self._parent(index)

    def find_min(self):
        """查看当前堆最小元素"""
        if self.is_empty():
            raise DataStructureIsEmpty("MaxHeap is Empty.")
        return self._data[0]

    def extract_min(self):
        """返回并删除堆中最小元素"""
        result = self.find_min()
        self._data[0], self._data[-1] = self._data[-1], self._data[0]
        self._data.pop()
        self.shift_down(0)
        return result

    def shift_down(self, index):
        """元素下沉，将较大元素放置到底层"""
        while self._left(index) < self.size():
            j = self._left(index)
            if j + 1 < self.size() and self._data[j + 1] < self._data[j]:
                j = self._right(index)
                # j为左右孩子中元素较小的节点的索引
            if self._data[index] <= self._data[j]:
                # 比左右孩子节点都大，不用下沉
                break
            self._data[index], self._data[j] = self._data[j], self._data[index]
            index = j

    def replace(self, e):
        """取出最大元素，并替换成新元素e"""
        result = self.find_min()
        self._data[0] = e
        self.shift_down(0)
        return result

    @classmethod
    def heapify(cls, arr):
        """将任意数组整理成最大堆，并返回最大堆"""
        array = list(arr)
        heap = MinHeap()
        heap._data = array
        i = heap._parent(heap.size())
        while i >= 0:
            heap.shift_down(i)
            i -= 1
        return heap
